import os
import json
from datetime import datetime
from typing import Any

class ChatHistoryManager:
    def __init__(self, base_data_dir="data"):
        self.base_data_dir = base_data_dir
        os.makedirs(self.base_data_dir, exist_ok=True)

    def _get_chat_history_path(self, userid: str) -> str:
        user_data_dir = os.path.join(self.base_data_dir, userid)
        os.makedirs(user_data_dir, exist_ok=True)
        return os.path.join(user_data_dir, "chat_history.txt")

    def append_chat_history(self, userid: str, user_content: Any, assistant_content: Any):
        history_path = self._get_chat_history_path(userid)
        
        all_messages = []
        if os.path.exists(history_path):
            try:
                with open(history_path, "r", encoding="utf-8") as f:
                    content = f.read()
                    if content:
                        all_messages = json.loads(content)
            except json.JSONDecodeError:
                all_messages = []

        current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

        user_entry = {
            "role": "user",
            "content": user_content,
            "datetime": current_time
        }
        assistant_entry = {
            "role": "assistant",
            "content": assistant_content,
            "datetime": current_time
        }

        all_messages.append(user_entry)
        all_messages.append(assistant_entry)

        with open(history_path, "w", encoding="utf-8") as f:
            json.dump(all_messages, f, ensure_ascii=False, indent=4)

    def get_recent_chat_history(self, userid: str, max_bytes: int = 61440) -> str:
        history_path = self._get_chat_history_path(userid)
        if not os.path.exists(history_path):
            return ""

        all_messages = []
        try:
            with open(history_path, "r", encoding="utf-8") as f:
                content = f.read()
                if content:
                    all_messages = json.loads(content)
        except json.JSONDecodeError as e:
            print(f"Error decoding JSON from chat history file: {e}")
            return ""
            
        recent_messages_formatted = []
        current_size = 0
        
        for msg in reversed(all_messages):
            formatted_msg = ""
            role = msg.get("role")
            content = msg.get("content", "")

            if isinstance(content, list):
                # Handle content as a list (e.g., for image URLs)
                content_str = ""
                for item in content:
                    if isinstance(item, dict) and "image_url" in item:
                        content_str += f"[image_url: {item['image_url'].get('url', '')}] "
                    else:
                        content_str += str(item) + " "
                content = content_str.strip()
            
            if role == "user":
                formatted_msg = f"user: {content}"
            elif role == "assistant":
                formatted_msg = f"assistant: {content}"
            
            if formatted_msg:
                msg_size = len(formatted_msg.encode("utf-8")) + (2 if recent_messages_formatted else 0)
                if current_size + msg_size <= max_bytes:
                    recent_messages_formatted.append(formatted_msg)
                    current_size += msg_size
                else:
                    break
        
        recent_messages_formatted.reverse()
        
        return "\n\n".join(recent_messages_formatted) + ("\n\n" if recent_messages_formatted else "")
